import numpy as np
import torch
import torchvision
import tensorboard
from torchvision import datasets, transforms
from matplotlib import pyplot as plt
from torch.utils.tensorboard import SummaryWriter

if __name__ == '__main__':

    transform = transforms.Compose([
        # transforms.Resize((224, 224)),  # 调整图像大小
        transforms.Resize((200, 200)),  # becasue vgg takes 150*150
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((.5, .5, .5), (.5, .5, .5))

    ])

    # 创建数据集
    data_dir = '/Users/xianda/Downloads/flower_photos/'  # 你的数据集路径
    batch_size = 32

    train_dataset = datasets.ImageFolder(data_dir, transform=transform)

    # 数据集拆分比例
    train_size = int(0.7 * len(train_dataset))
    val_size = len(train_dataset) - train_size

    # 随机拆分数据集
    train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])

    # 数据加载器
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    # default `log_dir` is "runs" - we'll be more specific here
    writer = SummaryWriter('../../runs/tensorBoard')

    dataiter = iter(train_loader)
    images, labels = next(dataiter)

    # create grid of images
    img_grid = torchvision.utils.make_grid(images)


    def matplotlib_imshow(img, one_channel=False):
        if one_channel:
            img = img.mean(dim=0)
        img = img / 2 + 0.5  # unnormalize
        npimg = img.numpy()
        if one_channel:
            plt.imshow(npimg, cmap="Greys")
        else:
            plt.imshow(np.transpose(npimg, (1, 2, 0)))


    # show images
    matplotlib_imshow(img_grid, one_channel=True)

    # write to tensorboard
    writer.add_image('four_fashion_mnist_images', img_grid)

    writer.close()